import torch
import json

inChanel = 3
outChanel = 1

input = torch.randn(inChanel)
weight = torch.randn(outChanel, inChanel)
bias = torch.randn(outChanel)
params = {'weight': weight, 'bias': bias}
operator = torch.nn.Linear(inChanel, outChanel)
operator.load_state_dict(params)
output = operator.forward(input)

exportData = [
    ('input', input),
    ('weight', weight),
    ('bias', bias),
    ('output', output)
]

data = []
for name, tensor in exportData:
    data.append({"layer": name, "value": tensor.tolist()})

jsonString = json.dumps(data)
with open('LinearTest.json', 'w') as f:
    f.write(jsonString)
print('TestBench generate finished.')
